package standalone.classification.linear;


import java.io.DataOutputStream;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

import utils.DataInput;
import utils.MathUtils;
import utils.StructureTransform;

import cern.colt.matrix.DoubleMatrix1D;
import cern.colt.matrix.DoubleMatrix2D;

/**
 * Hierarchical Logistic Regression 
 * @author purlin
 *
 */
public class HierarchicalLR {
	/**
	 * Node for the hierarchical structure
	 * @author purlin
	 *
	 */
	public class ModelTreeNode {
		/**
		 * Logistic Regression Object
		 */
		public LogisticRegression lrObj = null;
		/**
		 * Label of the parent
		 */
		public String parent = "";
		/**
		 * Labels of children
		 */
		public List<String> children = new ArrayList<String>();
	}
	
	/**
	 * Tree structure
	 */
	private Map<String, ModelTreeNode> modelTree = null;
	/**
	 * Labels of data
	 */
	private List<String> labels = null;
	/**
	 * Data, Sparse/Dense matrix
	 */
	private DoubleMatrix2D data = null;
	/**
	 * Number of features
	 */
	private int numFeatures = 0;
	/**
	 * Trade-off parameter
	 */
	private float d = 0.1f;
	/**
	 * Number of iterations for SGD
	 */
	private int numIt = 500;
	/**
	 * Learning rate
	 */
	private float learningRate = 0.01f;
	/**
	 * Termination threshold
	 */
	private float eps = 0.001f;
	/**
	 * IP address of spanning tree server
	 */
	protected String master = "localhost";
	/**
	 * Port
	 */
	protected int masterPort = 10090;
	
	/**
	 * Constructor
	 * @param numFeatures: number of features
	 * @param d: trade-off parameter
	 * @param numIt: number of iterations
	 * @param learningRate: learning rate
	 * @param eps: termination threshold
	 * @param master: master IP
	 * @param masterPort: master port
	 */
	public HierarchicalLR(int numFeatures, float d, int numIt, float learningRate, float eps, String master, int masterPort) {
		this.numFeatures = numFeatures;
		this.d = d;
		this.numIt = numIt;
		this.learningRate = learningRate;
		this.eps = eps;
		this.master = master;
		this.masterPort = masterPort;
	}
	
	/**
	 * Load data from disk
	 * @param inputPath: input path of data file
	 * @param fileType: file type: sparse(libsvm)/dense
	 * @throws IOException
	 */
	protected void loadData(String inputPath, String fileType) throws IOException{
		if("libsvm".endsWith(fileType)) data = DataInput.readHLibSVM(inputPath, numFeatures, labels);
		else if("dense".endsWith(fileType)) data = DataInput.readHDenseMatrix(inputPath, numFeatures, labels);
	}
	
	/**
	 * Build the hierarchical information of labels 
	 */
	public void buildHierarchical(){
		//Build a class-id map
		Set<String> classSet = new HashSet<String>();
		for(String label : labels){
			String[] items = label.split(".");
			String key = "";
			for(int i=0;i<items.length;i++){
				if(i!=0) key += ".";
				key = key + items[i];
				if(!classSet.contains(key)) classSet.add(key);
			}
		}
		//Build tree structure
		modelTree = new HashMap<String, ModelTreeNode>();
		String[] keys = classSet.toArray(new String[0]);
		for(int i=0;i<keys.length;i++){	//Initialization
			ModelTreeNode currentNode = new ModelTreeNode();
			modelTree.put(keys[i], currentNode);
		}
		for(int i=0;i<keys.length;i++){	//Parent-Children relation
			ModelTreeNode currentNode = modelTree.get(keys[i]);
			for(int j=0;j<keys.length;j++){
				if(i==j) continue;
				if(keys[j].startsWith(keys[i])) {
					ModelTreeNode otherNode = modelTree.get(keys[j]);
					currentNode.children.add(keys[j]);
					otherNode.parent = keys[i];
					modelTree.put(keys[j], otherNode);
				}
			}
			modelTree.put(keys[i], currentNode);
		}
		//Root
		ModelTreeNode root = new ModelTreeNode();
		for(int i=0;i<keys.length;i++){
			ModelTreeNode currentNode = modelTree.get(keys[i]);
			if(currentNode.parent.equals("")) root.children.add(keys[i]);
		}
		modelTree.put("root", root);
		//Initialize
		keys = modelTree.keySet().toArray(new String[0]);
		for(int i=0;i<keys.length;i++){
			ModelTreeNode currentNode = modelTree.get(keys[i]);
			modelTree.put(keys[i], currentNode);
		}		
	}
	
	/**
	 * Build model recursively
	 * @param label: label of the node
	 * @throws IOException
	 */
	protected void buildRecursion(String label) throws IOException{
		ModelTreeNode currentNode = modelTree.get(label);
		if(currentNode.children.size()<2) return;	//Leaf node
		//Generate current dataset
		List<Integer> currentLabels = new ArrayList<Integer>();
		List<DoubleMatrix1D> currentData = new ArrayList<DoubleMatrix1D>();
		for(int i=0;i<labels.size();i++){
			String tag = labels.get(i);
			for(int j=0;j<currentNode.children.size();j++){
				if(tag.startsWith(currentNode.children.get(j))) {
					currentLabels.add(j);
					currentData.add(data.viewRow(i));
				}
			}
		}
		//Model building
		currentNode.lrObj = new LogisticRegression(numFeatures, currentNode.children.size());
		currentNode.lrObj.initialize(master, masterPort);
		currentNode.lrObj.gradientDescent(StructureTransform.vecList2Matrix(currentData), currentLabels, d, numIt, learningRate, eps);
		modelTree.put(label, currentNode);
		//Sub-model building
		for(int i=0;i<currentNode.children.size();i++) buildRecursion(currentNode.children.get(i));
	}
	
	/**
	 * Predict an instance's label recursively 
	 * @param label: label of parent
	 * @param instance: given instance
	 * @return label: predicted label
	 */
	protected String predictRecursion(String label, DoubleMatrix1D instance){
		ModelTreeNode currentNode = modelTree.get(label);
		if(currentNode.children.size()<2) return label;
		String currentLabel = currentNode.children.get(MathUtils.getMaxID(currentNode.lrObj.predictProb(instance)));
		return predictRecursion(currentLabel, instance);
	}
	
	/**
	 * Save model recursively
	 * @param output: output stream
	 * @param label: label
	 * @throws IOException
	 */
	protected void saveRecursion(DataOutputStream output, String label) throws IOException{
		ModelTreeNode currentNode = modelTree.get(label);
		output.write((label+"\t"+currentNode.children.size()+"\n").getBytes());		
		for(int i=0;i<currentNode.children.size();i++)
			output.write((currentNode.children.get(i)+"\t"+currentNode.lrObj.w[i].toString()+"\n").getBytes());
		for(int i=0;i<currentNode.children.size();i++) saveRecursion(output, currentNode.children.get(i));		
	}
	
	/**
	 * Predict the labels of all instances
	 * @param testPath: path of testing data
	 * @param dataType: data format: libsvm or dense
	 * @param prePath: path of prediction
	 * @throws IOException
	 */
	public void predict(String testPath, String dataType, String prePath) throws IOException{
		labels.clear(); data = null;
		if("libsvm".endsWith(dataType)) data = DataInput.readHLibSVM(testPath, numFeatures, labels);
		else if("dense".endsWith(dataType)) data = DataInput.readHDenseMatrix(testPath, numFeatures, labels);
		DataOutputStream out = new DataOutputStream(new FileOutputStream(new File(prePath)));
		for(int i=0;i<data.rows();i++)
			out.write((predictRecursion("root",data.viewRow(i))+"\n").getBytes());
		out.close();
	}
	
	/**
	 * Model Building
	 * @param inputPath: training data path
	 * @param outputPath: output path
	 * @param dataType: data representation type: libsvm/dense
	 * @param modelName: built model name
	 * @param d: trade-off parameter
	 * @param numFeatures: number of features 
	 * @param numClasses: number of classes
	 * @param numIt: number of iterations
	 * @param learningRate: learning rate
	 * @param eps: threshold for termination 
	 * @param master: IP of spanning tree server
	 * @param masterPort: Port of spanning tree server
	 * @throws Exception
	 */
	public static void buildModel(String inputPath, String outputPath, String dataType, String testPath, String prePath, String modelName, float d, int numFeatures, int numClasses, int numIt, float learningRate, float eps, String master, int masterPort) throws NumberFormatException, IOException {
		HierarchicalLR model = new HierarchicalLR(numFeatures, d, numIt, learningRate, eps, master, masterPort);
		model.loadData(inputPath, dataType);
		model.buildHierarchical();
		model.buildRecursion("root");
		DataOutputStream out = new DataOutputStream(new FileOutputStream(new File(outputPath+"/"+modelName)));
		model.saveRecursion(out, "root");
		out.close();
		model.predict(testPath, dataType, prePath);
	}
	
}
